import cv2
import numpy as np
import math
import os

MOON_RADIUS = 1737.4  # km
BATCH_SIZE = 1000
## for M1348581418LE_DEM
# LATITUDE_RANGE = 1.935  # degree
# LONGITUDE_RANGE = 0.5  # degree
# X_PIXEL_RANGE = 58235
# Y_PIXEL_RANGE = 15042


# def simple_cylinder(lat, lon, lat0=44.43, lon0=-52.24):
#     x = (lat0 - lat) * X_PIXEL_RANGE / LATITUDE_RANGE
#     y = (lon - lon0) * Y_PIXEL_RANGE / LONGITUDE_RANGE
#     return x, y

## for M1173414625_DEM
LATITUDE_RANGE = 45.2410 - 42.8613  # degree
LONGITUDE_RANGE = -51.6879 - (-52.0974)  # degree
X_PIXEL_RANGE = 48131
Y_PIXEL_RANGE = 8264


def simple_cylinder(lat, lon, lat0=45.2410, lon0=-52.0974):
    x = (lat0 - lat) * X_PIXEL_RANGE / LATITUDE_RANGE
    y = (lon - lon0) * Y_PIXEL_RANGE / LONGITUDE_RANGE
    return x, y


def simple_cylinder_R(d, lat):
    D_x = math.degrees(d / MOON_RADIUS) * X_PIXEL_RANGE / LATITUDE_RANGE
    D_y = (
        math.degrees(d / MOON_RADIUS)
        * Y_PIXEL_RANGE
        / LONGITUDE_RANGE
        / math.cos(math.radians(lat))
    )

    return D_x / 2, D_y / 2


def labels(lat, lon, d):
    temp_1, temp_2 = simple_cylinder(lat, lon)
    R_x, R_y = simple_cylinder_R(d, lat)
    x_folder = int(temp_1) // BATCH_SIZE
    y_folder = int(temp_2) // BATCH_SIZE
    x = temp_1 - BATCH_SIZE * x_folder
    y = temp_2 - BATCH_SIZE * y_folder
    cir_labels = {"Y_FOLDER": y_folder, "X_FOLDER": x_folder, "Y_SUB": y, "X_SUB": x}

    return cir_labels, temp_1, temp_2, R_x, R_y


if __name__ == "__main__":
    FILE_PATH = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/regularized_M1173414625_DEM.tif"
    # LABEL_DATA_PATH = (
    #     "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/Qian_chang_e_5_annotation.csv"
    # )
    LABEL_DATA_PATH = (
        "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/Liu_chang_e_5_annotation.txt"
    )

    SUBFIG_DIR = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"
    # img = cv2.imread(FILE_PATH, cv2.IMREAD_UNCHANGED)
    # 从标注文件中获取某一行
    with open(LABEL_DATA_PATH, "r") as f:
        # 有部分行只有圆的数据，没有椭圆的数据
        # 流式读取，第一行是标题，并使用标题的字段名作为数据变量名
        header = f.readline()
        header = header.strip().split(",")
        for i, line_ in enumerate(f):
            line = line_.strip().split(",")
            name = line[0]
            lon = float(line[1])
            lat = float(line[2])
            d = float(line[3])
            # name = i
            # lon = float(line[0]) - 360
            # lat = float(line[1])
            # d = float(line[2]) / 1000
            cir_labels, temp1, temp2, Rx, Ry = labels(lat, lon, d)
            cir_x_folder = cir_labels["X_FOLDER"]
            cir_y_folder = cir_labels["Y_FOLDER"]
            x = cir_labels["X_SUB"]
            y = cir_labels["Y_SUB"]
            if (
                cir_x_folder > 48
                or cir_x_folder < 0
                or cir_y_folder > 8
                or cir_y_folder < 0
            ):
                continue
            if not os.path.exists(f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}"):
                os.makedirs(f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}")
            with open(
                f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}/chang_e_{cir_x_folder}_{cir_y_folder}_{BATCH_SIZE}.txt",
                "a",
            ) as writer:
                writer.write(
                    f"{x},{y},{cir_x_folder},{cir_y_folder},{name},{lat},{lon},{d}\n"
                )
                writer.flush()
    # for i in range(0, Y_PIXEL_RANGE, BATCH_SIZE):
    #     for j in range(0, X_PIXEL_RANGE, BATCH_SIZE):
    #         sub_img = img[i : i + BATCH_SIZE, j : j + BATCH_SIZE, :]
    #         sub_img = (sub_img - sub_img.min()) / (sub_img.max() - sub_img.min()) * 255
    #         cv2.imwrite(f"test_{i//BATCH_SIZE}_{j//BATCH_SIZE}.png", sub_img)
